package com.btl.service.ai.biz.adaptor;

import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.btl.boot.core.dto.response.R;
import com.btl.service.ai.auth.SecurityContextHolder;
import com.btl.service.ai.biz.dto.segment.ModSegmentQuery;
import com.btl.service.ai.biz.entity.KwSegment;
import com.btl.service.ai.biz.service.IKwSegmentService;
import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingModelTypeEnum;
import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingStoreIngestorStrategy;
import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingStoreIngestorStrategyFactory;
import com.btl.service.ai.support.rag.store.vector.collections.VectorCollectionsEnum;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.Metadata;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.annotation.Resource;
import jakarta.validation.Valid;
import jakarta.validation.constraints.NotNull;
import lombok.extern.slf4j.Slf4j;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.StringUtils;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;
import org.springframework.stereotype.Controller;

import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * <p>
 * 前端控制器
 * </p>
 *
 * @author Tim
 * @since 2025-07-22
 */
@Slf4j
@Validated
@RestController
@Tag(name = "分段")
@Controller
@RequestMapping("segment")
public class KwSegmentController {

    @Resource
    private IKwSegmentService iKwSegmentService;

    /**
     * 分段分页
     *
     * @param current    当前页
     * @param size       每页数量
     * @param documentId 文档ID
     * @return R<Page < KwDocument>>
     */
    @GetMapping("page/{documentId}")
    public R<Page<KwSegment>> page(
            @PathVariable Long documentId,
            @RequestParam(value = "current", required = false) @NotNull(message = "当前页码不能为空") Integer current,
            @RequestParam(value = "size", required = false) @NotNull(message = "每页数量不能为空") Integer size) {
        return R.ok(iKwSegmentService.page(documentId, current, size));
    }

    /**
     * 分段修改
     *
     * @param modSegmentQuery 请求信息
     * @return R<Page < KwDocument>>
     */
    @PutMapping("mod/{segmentId}")
    @Transactional(rollbackFor = Exception.class)
    public R<Map<String, Object>> mod(@PathVariable Long segmentId, @RequestBody @Valid ModSegmentQuery modSegmentQuery) {
        Long childId = SecurityContextHolder.get()
                .getChildId();
        KwSegment kwSegment = iKwSegmentService.getById(segmentId);
        if (!Objects.equals(kwSegment.getChildId(), childId)) {
            return R.fail("分段不存在");
        }
        String title = Objects.isNull(kwSegment.getTitle()) ? "" : kwSegment.getTitle();
        iKwSegmentService.updateById(KwSegment.builder()
                .segment(modSegmentQuery.getSegment())
                .title(title)
                .id(segmentId)
                .build());

        String segment = kwSegment.getSegment();
        if (StringUtils.hasText(title)) {
            segment = title + "\n\n" + segment;
        }
        Document document = Document.from(segment, new Metadata(Map.of(
                "vectorId", kwSegment.getVectorId(),
                "index", kwSegment.getSeqId(),
                "segmentId", kwSegment.getId()
                        .toString(),
                "documentId", kwSegment.getDocumentId()
                        .toString()
        )));
        String collectionName = VectorCollectionsEnum.KW_SEGMENT.getCollection();
        EmbeddingStoreIngestorStrategy embeddingStoreIngestorStrategy = EmbeddingStoreIngestorStrategyFactory.ingestorStrategy(EmbeddingModelTypeEnum.LOCAL_BGE_SMALL_ZH_V15_EMBEDDING_MODEL);
        embeddingStoreIngestorStrategy.embed(collectionName, List.of(document));
        return R.ok(Map.of("segmentId", segmentId));
    }
}
